
from os.path import join
from models import KDAD
from torch.utils.data import DataLoader
from datasets import mvtec


def kdad_run(cfg,phase,weights=''):

    print('using kdad')

    category = cfg['normal_class']
    train_batch_size = cfg['train_batch_size']
    test_batch_size = cfg['test_batch_size']
    load_size = cfg['load_size']
    input_size = cfg['input_size']

    if category == 'all':
        train_class = mvtec.CLASS_NAMES
    else:
        train_class = [category]

    model = KDAD(cfg)
    # print(train)
    for c in train_class:
        print(c)
        train_dataset = mvtec.MVTecDataset(
            root_path=cfg['dataset_dir'], class_name=c, is_train=True, resize=load_size, cropsize=input_size)
        train_dataloader = DataLoader(
            train_dataset, batch_size=train_batch_size, pin_memory=True, num_workers=0)
        test_dataset = mvtec.MVTecDataset(
            root_path=cfg['dataset_dir'], class_name=c, is_train=False, resize=load_size, cropsize=input_size)
        test_dataloader = DataLoader(
            test_dataset, batch_size=test_batch_size, pin_memory=True, num_workers=0)

        if phase == 'train':
            model.train(train_dataloader, test_dataloader)

        elif phase == 'test':
            model.test_after(test_dataloader, c)
            model.evaluate(c)
